"""
编码器
"""
import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

import config


class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.embedding = nn.Embedding(num_embeddings=len(config.chatbot_ws_input_by_word_model),
                                      embedding_dim=config.chatbot_embedding_dim,
                                      padding_idx=config.chatbot_ws_input_by_word_model.PAD)
        self.gru = nn.GRU(input_size=config.chatbot_embedding_dim,
                          hidden_size=config.chatbot_encoder_hidden_size,
                          num_layers=config.chatbot_encoder_num_layers,
                          batch_first=True)

    def forward(self, input_data, input_length):
        # input_data: [batch_size, max_len]
        # input_length: [batch_size,]

        # 1. Embedding操作
        # embed: [ batch_size, max_len, embedding_dim]
        embed = self.embedding(input_data)

        # 2. 打包
        # output: [batch_size, max_len, embedding_dim]
        # print("input_length:", input_length.size())
        output = pack_padded_sequence(embed, input_length, batch_first=True, enforce_sorted=True)

        # 3. GRU
        # output: [batch_size, max_len, hidden_size]
        # hidden_state: [1, batch_size, hidden_size]
        output, hidden_state = self.gru(output)
        # print("output:", output.size())
        # print("hidden_state:", hidden_state.size())

        # 4. 解包
        output, output_length = \
            pad_packed_sequence(output, batch_first=True, padding_value=
                                config.chatbot_ws_input_by_word_model.PAD)

        # print("output:", output.size())
        return output, output_length, hidden_state





